#
# Copyright (C) 2023, Inria
# GRAPHDECO research group, https://team.inria.fr/graphdeco
# All rights reserved.
#
# This software is free for non-commercial, research and evaluation use 
# under the terms of the LICENSE.md file.
#
# For inquiries contact  george.drettakis@inria.fr
#

import torch
from kornia import create_meshgrid
from torch.fft import fft2, fftshift, ifft2, ifftshift

def mse(img1, img2):
    return (((img1 - img2)) ** 2).view(img1.shape[0], -1).mean(1, keepdim=True)

def psnr(img1, img2):
    mse = (((img1 - img2)) ** 2).view(img1.shape[0], -1).mean(1, keepdim=True)
    return 20 * torch.log10(1.0 / torch.sqrt(mse))

def distance_from_center_to_coord(coords):
        return torch.sqrt(coords[..., 0]**2 + coords[..., 1]**2)
    
def get_mask(H, W, r, buffer_width):
    R = r + buffer_width
    mask = torch.ones((H, W))
    grid = grid = create_meshgrid(H, W, normalized_coordinates=False)[0]
    i, j = grid.unbind(-1)
    coord = torch.stack([i-W//2+1, j-H//2+1], -1).abs()
    dis = distance_from_center_to_coord(coord)

    mask[dis > r] = (R - dis[dis > r]) / (R - r)
    mask[dis > R] = 0.0
    
    return mask

def img2fft(X):
    return fftshift(fft2(X))

def fft2img(X):
    return ifft2(ifftshift(X)).real

def soft_linear_split_freq(X, r:int, buffer_width:float=5.0):
    """
    ## Parameters
        - X: Tensor in shape [C, H, W]. Input image.
        - r: int, radii of freq map defined low freq.
        - buffer_width: buffer width from low to high
    ## Returns
        - X_freqs: List[Tensor]. The splitted images of each freq range.
    """
    
    H, W = X.shape[-2:]
    mask = get_mask(H, W, r, buffer_width).to(X.device)
    freq = img2fft(X)
    low_freq = freq * mask
    
    low_freq_img = fft2img(low_freq)
    high_freq_img = X - low_freq_img
    
    return [low_freq_img, high_freq_img]